#
# regression.py
#
# All Copyright (C) reserved.
# shenyczz@163.com
# ----------------------------------------------------------
#
import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt

EPOCH = 150


class CRegression(torch.nn.Module):
    '''
    回归:  y=x^3 + b
    '''

    def __init__(self, n_feature, n_hidden, n_output):
        super(CRegression, self).__init__()
        self.hidden = torch.nn.Linear(n_feature, n_hidden)
        self.predict = torch.nn.Linear(n_hidden, n_output)

    def forward(self, x):
        x = F.relu(self.hidden(x))
        return self.predict(x)


def regression():
    # 准备数据: y = x^3 + b
    x = torch.unsqueeze(torch.linspace(-5, 5, 100), dim=1)
    y = x.pow(3) + 0.2 * torch.rand(x.size())

    # 确定输入层、隐藏层、输出层的神经元数目
    net = CRegression(1, 20, 1)
    print(net)
    optimizer = torch.optim.Adam(net.parameters(), lr=0.2)
    loss_func = torch.nn.MSELoss()

    plt.ion()

    for epoch in range(EPOCH):
        prediction: torch.Tensor = net(x)
        loss: torch.Tensor = loss_func(prediction, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if epoch % 5 == 0:
            plt.cla()
            plt.scatter(x.data.numpy(), y.data.numpy(), s=10, c='b')
            plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=2)
            plt.text(-2, 100, 'Epoch=%i Loss=%.4f' % (epoch, loss.data.numpy()),
                     fontdict={'size': 10, 'color': 'red'})
            # plt.savefig(f'e:/temp/{epoch}.png')
            plt.pause(0.5)
        # if epoch % 5 == 0
    # for t in range(2000)

    plt.ioff()
    plt.show()
# regression():


regression()
